import datetime

from vtkmodules.vtkRenderingCore import (
    vtkShaderProperty,
    vtkTexture,
    vtkSkybox,
    vtkActor,
    vtkTextActor,
    vtkPolyDataMapper,    
    vtkProperty,
    vtkCamera,    
    vtkLight,
    vtkHardwareSelector
)

from vtkmodules.vtkInteractionWidgets import vtkTextWidget

from vtkmodules.vtkRenderingOpenGL2 import (
    vtkOpenGLRenderer,
    vtkRenderStepsPass,
    vtkCameraPass,
    vtkOrderIndependentTranslucentPass,
    vtkSSAOPass,    
    vtkRenderPassCollection,
    vtkSequencePass,
    vtkShadowMapPass,
    vtkOutlineGlowPass,
    vtkOpenGLState
)

from vtkmodules.vtkIOImage import vtkHDRReader

from vtkmodules.vtkFiltersCore import vtkAppendPolyData
from vtkmodules.vtkCommonTransforms import vtkTransform


from vtkmodules.vtkCommonCore import (
    vtkPoints,
)

from vtkmodules.vtkCommonDataModel import (
    vtkPolyData,
    vtkCellArray,
)

from vtkmodules.vtkFiltersSources import (
    vtkPlaneSource, 
    vtkDiskSource,
    vtkCylinderSource
)

from vtkmodules.vtkInteractionWidgets import (
    vtkCameraOrientationWidget,
    vtkLightWidget,
    vtkLightRepresentation,
    vtkBoxWidget2,
    vtkBoxRepresentation
)

from vtkmodules.vtkCommonDataModel import vtkSelectionNode


class GLTFRenderer(vtkOpenGLRenderer):

    def __init__(self) -> None:
        super().__init__()         
        #self.SetUseDepthPeeling(1)
        #self.SetMaximumNumberOfPeels(16)


        # setup render pass
        # basicPasses = vtkRenderStepsPass()   
        # # opacity
        # peeling = vtkOrderIndependentTranslucentPass()
        # peeling.SetTranslucentPass(basicPasses.GetTranslucentPass())
        # basicPasses.SetTranslucentPass(peeling)
        # # ssao pass
        # ssao = vtkSSAOPass()
        # ssao.SetRadius(0.05)
        # ssao.SetKernelSize(128)
        # ssao.SetDelegatePass(basicPasses)  
        # cameraPass = vtkCameraPass()
        # cameraPass.SetDelegatePass(basicPasses)
        # self.SetPass(ssao)
        # shadows = vtkShadowMapPass()
        # seq = vtkSequencePass()
        # passes = vtkRenderPassCollection()
        # passes.AddItem(shadows.GetShadowMapBakerPass())
        # passes.AddItem(shadows)
        # seq.SetPasses(passes)
        # cameraP = vtkCameraPass()
        # cameraP.SetDelegatePass(seq)
        #self.SetPass(cameraP)
 
        # show framerate           
        self.showFPS = False

        self.fpsActor = vtkTextActor()
        self.fpsActor.SetPosition(20, 20)
        self.fpsActor.GetTextProperty().SetFontSize(18)
        self.AddActor2D(self.fpsActor)                

        self.frameStart = None
        self.frameCount = 0

        def fpsCallback(caller, evt):
            if self.showFPS:
                if self.frameStart == None:
                    self.frameStart = datetime.datetime.now()
                    self.frameCount = 0
                    return
                
                self.frameCount += 1
                now = datetime.datetime.now()
                elapsed = (now - self.frameStart).total_seconds()

                if elapsed > 1:
                    self.fpsActor.SetInput("{:.0f} ms".format(elapsed * 1000 / self.frameCount))                    
                    self.frameStart = now
                    self.frameCount = 0
           
        self.AddObserver("EndEvent", fpsCallback)

        self.groundActor = None
        self.skybox = None

        self.sceneActors = []

    def SetEnvmap(self, path) -> vtkSkybox:

        if self.skybox == None:
            self.skybox = vtkSkybox()    
            self.skybox.SetFloorRight(1, 0, 0)
            self.skybox.SetProjection(vtkSkybox.Sphere)
            self.skybox.GammaCorrectOn()
            self.AddActor(self.skybox)

        if path:
            # load IBL
            reader = vtkHDRReader()
            reader.SetFileName(path)
            reader.Update()

            texture = vtkTexture()
            texture.MipmapOn()
            texture.InterpolateOn()
            texture.SetColorModeToDirectScalars()
            texture.SetInputConnection(reader.GetOutputPort())

            self.UseImageBasedLightingOn()
            self.UseSphericalHarmonicsOn()
            self.SetEnvironmentTexture(texture)
            self.skybox.SetTexture(texture)
            self.skybox.SetVisibility(1)
        else:            
            self.UseImageBasedLightingOff()
            self.UseSphericalHarmonicsOff()
            self.skybox.SetVisibility(0)

        return self.skybox
    
    def SetGround(self, on:int, center=[0, 0, 0], scale=1.0, upAxis='y'):
        
        if self.groundActor == None:
            self.groundActor = vtkActor()
            self.groundActor.GetProperty().SetInterpolationToPBR()
            self.groundActor.GetProperty().SetColor(0.1, 0.1, 0.1)
            self.groundActor.GetProperty().SetMetallic(0)
            self.groundActor.GetProperty().SetRoughness(0.3)
            self.AddActor(self.groundActor)
            
            # create a ground
            source = vtkDiskSource()
            source.SetInnerRadius(0)
            source.SetOuterRadius(1)
            source.SetCircumferentialResolution(36)
            source.Update()
            mapper = vtkPolyDataMapper()
            mapper.SetInputConnection(source.GetOutputPort())
            self.groundActor.SetMapper(mapper)

        if on:
            if upAxis == 'y':
                self.groundActor.RotateX(90)
            self.groundActor.SetScale(scale)
            self.groundActor.SetPosition(center)
            self.groundActor.SetVisibility(1)
        else:
            self.groundActor.SetVisibility(0)


    def AddScene(self, scene):        
        # add actors
        scene.InitTraversal()
        for actor in scene:
            self.AddActor(actor)            
            self.sceneActors.append(actor)

    def ClearScene(self):
        for actor in self.sceneActors:
            self.RemoveActor(actor)
        self.sceneActors.clear()

# draw outline for target scene node
class SceneNodeOutlineRenderer(vtkOpenGLRenderer):

    def __init__(self):        
        # setup render pass
        basicPasses = vtkRenderStepsPass()
        glowPass = vtkOutlineGlowPass()
        glowPass.SetDelegatePass(basicPasses)
        glowPass.SetOutlineIntensity(5.0) 
        cameraPass = vtkCameraPass()
        cameraPass.SetDelegatePass(glowPass)
        # Apply the render pass to the highlight renderer
        self.SetPass(cameraPass)

        self.property = vtkProperty()
        self.property.LightingOff()
        self.property.SetOpacity(0.8)
        self.property.SetColor(1,1,0)

    def AddSceneNode(self, scene, node):
        # create a actor for outline
        polyData = vtkAppendPolyData()
        polyData.ParallelStreamingOn()
        for mesh in node.meshes:
            for prim in mesh.primitives:
                polyData.AddInputData(prim.GetOutput())
        
        actor = vtkActor()
        actor.SetProperty(self.property)
        
        transform = vtkTransform()
        transform.Concatenate(scene.transform)
        transform.Concatenate(node)
        actor.SetUserTransform(transform)
        
        mapper = vtkPolyDataMapper()
        mapper.SetInputConnection(polyData.GetOutputPort())
        actor.SetMapper(mapper)
        self.AddActor(actor)

    def Clear(self):
        self.RemoveAllViewProps()

class SceneNodeHighlightRenderer(vtkOpenGLRenderer):
    def __init__(self) -> None:
        super().__init__()     
        #self.SetPreserveColorBuffer(1)  
        #self.SetPreserveDepthBuffer(1)         
        # self.property = vtkProperty()
        # self.property.ShadingOff()
        # self.property.LightingOff()
        #self.property.SetOpacity(0.2)
        #self.property.SetColor(1,0,0)
        self.UseDepthPeelingOff()
        #self.SetMaximumNumberOfPeels(0)
        #self.property.BackfaceCullingOff()
        #self.property.FrontfaceCullingOff()
        #self.shaderProperty = vtkShaderProperty()
        # self.shaderProperty.AddFragmentShaderReplacement(
        #     "//VTK::Light::Impl",  True,
        #     "//VTK::Light::Impl\n"         
        #     "fragOutput0 = vec4(1, 0, 0, 0.2);\n"
        #     #"gl_FragDepth = gl_FragCoord.z - 0.00001;"
        #     , False)

        #def frameBegin(caller, evt):
        #    state = self.GetState()
        #    state.Push()
            #state.vtkglDisable(0x0B71) # GL_DEPTH_TEST
            #state.vtkglDisable(0x0B44) # GL_CULL_FACE
            #state.vtkglCullFace(0)
            #state.vtkglBlendEquation(0x8008) # GL_MAX
            #state.vtkglBlendEquationSeparate(0x8006, 0x8007) # GL_FUNC_ADD, GL_MIN
            #state.vtkglDepthMask(0)
            #state.vtkglDepthFunc(0x0203)
            #state.vtkglDepthFunc(0x0207) # always
            #state.vtkglBlendEquationSeparate(0x8006, 0x8006) # GL_FUNC_ADD, GL_MAX
            #state.vtkglDisable(0x0BE2) # GL_BLEND
            #state.ResetGLBlendFuncState()
            #state.ResetGLBlendEquationState()
            #state.vtkglBlendFunc(0, 0)
            #print("frameBegin")            
        # def frameEnd(caller, evt):
        #     #print("frameEnd")
        #     state = self.GetState()
        #     state.Pop()        
        # self.AddObserver("StartEvent", frameBegin)        
        # self.AddObserver("EndEvent", frameEnd)

    def AddSceneNode(self, scene, node, color, opacity):
        polyData = vtkAppendPolyData()
        polyData.ParallelStreamingOn()
        for mesh in node.meshes:
            for prim in mesh.primitives:
                polyData.AddInputData(prim.GetOutput())
        
        actor = vtkActor()
        actor.GetProperty().ShadingOff()
        actor.GetProperty().SetColor(color)
        actor.GetProperty().SetOpacity(opacity)
        #actor.SetShaderProperty(self.shaderProperty)    
        transform = vtkTransform()
        transform.Concatenate(scene.transform)
        transform.Concatenate(node)
        actor.SetUserTransform(transform)        
        mapper = vtkPolyDataMapper()
        mapper.SetInputConnection(polyData.GetOutputPort())
        actor.SetMapper(mapper)
        self.AddActor(actor)

    def Clear(self):
        self.RemoveAllViewProps()


class WidgetRenderer(vtkOpenGLRenderer):
    def __init__(self) -> None:
        super().__init__()

        # light control
        self.light = None
        # light representation 
        self.lightRep = vtkLightRepresentation()
        self.lightRep.SetPositional(1)

        def updateLight(caller, evt):
            if self.light:
                self.light.SetColor(self.lightRep.GetLightColor())
                self.light.SetPosition(self.lightRep.GetLightPosition())
                self.light.SetFocalPoint(self.lightRep.GetFocalPoint())
                self.light.SetConeAngle(self.lightRep.GetConeAngle())
                self.light.SetPositional(self.lightRep.GetPositional())
                #print("Light Position:", self.light.GetPosition())

        self.lightWidget = vtkLightWidget()
        self.lightWidget.SetRepresentation(self.lightRep)
        self.lightWidget.AddObserver("InteractionEvent", updateLight)
        self.lightWidget.SetDefaultRenderer(self)

        # add a camera axis widget
        self.cameraOrientationWidget = None
        #self.cameraOrientationWidget = vtkCameraOrientationWidget()

        # box for object manipulation
        self.boxRep = vtkBoxRepresentation()
        #self.boxRep.OutlineFaceWiresOn()
        self.boxRep.OutlineCursorWiresOff()
        self.boxRep.HandlesOff()
        self.boxRep.SetUseBounds(0)
        
        self.boxWidget = vtkBoxWidget2()
        self.boxWidget.SetRepresentation(self.boxRep)
        self.boxWidget.SetCurrentRenderer(self)
        self.boxWidget.MoveFacesEnabledOff()

        def updateTargetTransform(obj, evt):
            if self.boxWidgetTransform:
                self.boxRep.GetTransform(self.boxWidgetTransform)  
                self.ResetCameraClippingRange()   
        self.boxWidget.AddObserver("InteractionEvent", updateTargetTransform)

    def ToggleLightWidget(self):
        self.lightWidget.SetEnabled(not self.lightWidget.GetEnabled())            


    def SetInteractor(self, iren):    
        if self.cameraOrientationWidget:        
            self.cameraOrientationWidget.SetParentRenderer(self)
            self.cameraOrientationWidget.SetInteractor(iren)
            self.cameraOrientationWidget.On()
            # reset viewport of widget
            self.cameraOrientationWidget.GetDefaultRenderer().SetViewport(0.90, 0.90, 0.05, 0.05)
            self.cameraOrientationWidget.SquareResize()

        self.lightWidget.SetInteractor(iren)
        self.boxWidget.SetInteractor(iren)


    def SetBoxWidgetTransform(self, transform):
        self.boxWidgetTransform = transform
        if self.boxWidget.GetEnabled():
            self.boxWidget.SetEnabled(0)
        elif transform:
            self.boxRep.SetTransform(transform)
            self.boxWidget.SetEnabled(1)

    def SetLight(self, light):
        self.light = light
        if self.light:
            self.lightRep.SetLightPosition(self.light.GetPosition())
            self.lightRep.SetFocalPoint(self.light.GetFocalPoint())
            self.lightRep.SetLightColor(self.light.GetDiffuseColor())


# pick scene-node
class SceneNodePicker(vtkHardwareSelector):
    def __init__(self) -> None:
        pass
        #self.SetFieldAssociation(vtkDataObject.FIELD_ASSOCIATION_CELLS)

    def Pick(self, x, y):
        self.SetArea(x, y, x, y)
        res = self.Select()
        if res.GetNumberOfNodes() > 0:
            # only get the first one...
            node = res.GetNode(0)
            prop = node.GetProperties()
            id = node.GetSelectionList().GetValue(0) 

            key = vtkSelectionNode.PROP()
            if prop.Has(key):
                actor = key.Get(prop)
                # check if it is an instance actor...
                if hasattr(actor, 'scene'):
                    #return actor, id
                    scene = actor.scene
                    node  = actor.GetNode(id) 
                    return scene, node
                
        return None, None
    
# a set of renderers
class GLTFRenderEngine():
    def __init__(self) -> None:
        self.camera = vtkCamera()
        self.renderer = GLTFRenderer()
        self.renderer.SetActiveCamera(self.camera)
        self.outline = SceneNodeOutlineRenderer()
        self.outline.SetActiveCamera(self.camera)
        self.highlight = SceneNodeHighlightRenderer()
        self.highlight.SetActiveCamera(self.camera)

        # render widgets...
        self.widgets = WidgetRenderer()        
        self.widgets.SetActiveCamera(self.camera)
        #self.widgets.SetPreserveColorBuffer(1)
        #self.widgets.SetPreserveDepthBuffer(1)

        # add a default light
        self.light = vtkLight()
        self.light.SetPositional(1)
        self.light.SwitchOn()
        self.light.SetPosition(5, 5, -5)
        self.light.SetDiffuseColor(8, 8, 8)
        self.light.SetAmbientColor(0, 0, 0)

        self.renderer.AddLight(self.light)
        self.widgets.SetLight(self.light)

        self.picker = SceneNodePicker()
        self.picker.SetRenderer(self.renderer)

        # create another light
        self.renderer.CreateLight()

        # reset camera clip range before render
        def func(obj, evt):
            self.renderer.ResetCameraClippingRange()
        self.renderer.AddObserver("StartEvent", func)

    def GetCamera(self):
        return self.camera

    def GetDefaultRenderer(self):
        return self.renderer
    
    def GetOutlineRenderer(self):
        return self.outline
    
    def GetHighlightRenderer(self):
        return self.highlight
    
    def SetRenderWindow(self, window, layers:bool = True):  
        window.AddRenderer(self.renderer)
        window.AddRenderer(self.highlight)
        window.AddRenderer(self.outline)
        window.AddRenderer(self.widgets)

        if layers:
            window.SetNumberOfLayers(4)
            self.renderer.SetLayer(0)
            self.highlight.SetLayer(1)
            self.outline.SetLayer(2)
            self.widgets.SetLayer(3)
        else:
            self.highlight.SetPreserveColorBuffer(1)
            self.outline.SetPreserveColorBuffer(1)
            self.widgets.SetPreserveColorBuffer(1)        

    def SetInteractor(self, iren):
        self.widgets.SetInteractor(iren)

    def Pick(self, x, y):        
        scene, node = self.picker.Pick(x, y)
        if scene == None or node == None:            
            self.outline.Clear()
        else:    
            self.outline.Clear()
            self.outline.AddSceneNode(scene, node)
        return scene, node    

    def Highlight(self, scene, node, color=[1,0,0], opacity=1):
        if scene == None:
            return
        if node == None:            
            for node in scene.nodes:
                if len(node.meshes) > 0:
                    self.highlight.AddSceneNode(scene, node, color, opacity)
        else:
            self.highlight.AddSceneNode(scene, node, color, opacity)
        
    def ClearHighlight(self):
        self.highlight.Clear()



               
        
